#include<iostream>
#include<iomanip>
#include"sm_solver.h"

namespace SM
{
  void SMSolver::diag(int MM,int Par,DataType Emax)
  {
    if(world.size()!=1)
      {
	std::cerr<<"!!!Direct diagonalization requires only 1 process!\n";
	environment::abort(0);
      }
    if(world.rank()==0)
      {
	configs.setup(MM,Par,Emax);
	int dim=configs.dim;
	std::cout<<"total dim:\t"<<dim<<std::endl;
	numStates=dim;	
	if(dim==0)
	  {
	    eigVals.resize(0);
	    eigVecs.resize(0,0);
	    J2s.resize(0);
	  }
	else
	  {
	    Mat mat(dim,dim);
	    for(int i=0;i<dim;i++)
	      for(int j=i;j<dim;j++)
		{
		  mat(i,j)=configs.getME(i,j);
		  mat(j,i)=mat(i,j);
		}
	    ComplexEigenSolver<Mat> es(mat);
	    //	    SelfAdjointEigenSolver<Mat> es(mat);
	    eigVals=es.eigenvalues();
	    eigVecs=es.eigenvectors();

	    //norm
	    for(int i=0;i<dim;i++)
	      {
	    	DataType norm=eigVecs.col(i).transpose()*eigVecs.col(i);
	    	eigVecs.col(i)/=sqrt(norm);
	      }

	    //    sort
	    quickSort(eigVals,eigVecs,0,dim-1);

	    for(int i=0;i<dim;i++)
	      for(int j=i;j<dim;j++)
		{
		  mat(i,j)=configs.getJ2(i,j);
		  mat(j,i)=mat(i,j);
		}
	    J2s=(eigVecs.transpose()*mat*eigVecs).diagonal();
	  }
      }
    
    //needed for calOccuNums
    numOfparts=1;
    partdim=eigVals.size();
    diagcom=world;
  }


  
  void SMSolver::diag_lanczos(int num,int MM,int Par,DataType Emax,int step)
  {
    int rank=world.rank();

    if(rank==0)
      configs.setup(MM,Par,Emax);
    broadcast(world,configs,0);
  
    totaldim=configs.dim;
    if(rank==0)
      std::cout<<"total dim:\t"<<totaldim<<std::endl;
    if(totaldim==0)
      {
	eigVecs.resize(0,0);
	eigVals.resize(0);
	J2s.resize(0);
	numStates=0;
	numOfparts=1;	
	return;
      }

    if(num>totaldim) num=totaldim;
    numStates=num;
    
    LanczosIter<Configs> lanczositer(configs,world);
    
    int lanczosdim=num+step;

    double sigma;
    Vec eigVals_last;
    lanczositer.compute(eigVals,lanczosdim-1,false);    
    do
      {
    	eigVals_last=eigVals;
    	lanczositer.compute(eigVals,lanczosdim,false);

    	if(rank==0)
    	  {
    	    sigma=(eigVals.head(num)-eigVals_last.head(num)).norm();
    	    //show convergance
	    std::cout<<lanczosdim<<"\t"<<eigVals.head(num).transpose()<<"\t"<<sigma<<std::endl;
    	  }
    	broadcast(world,sigma,0);
    	lanczosdim+=step;
      }while(sigma>1e-8);

    

    lanczositer.compute(eigVals,lanczosdim,true);
    lanczositer.calJ2(J2s,numStates);

    //needed for calOccuNums
    lanczositer.getEigVecs(eigVecs,numStates);
    numOfparts=lanczositer.numOfparts;
    partdim=lanczositer.partdim;
    diagcom=lanczositer.diagcom;
  }
  
  void SMSolver::calOccuNums(int orb)
  {
    int rank=world.rank();
    if(rank<numOfparts)
      {
	int totalOrbitals=configs.Orbitals.size();
	OccuNums.resize(totalOrbitals);
	for(int i=0;i<totalOrbitals;i++)
	  {
	    OccuNums[i]=0;
	  }

	for(int j=0;j<totalOrbitals;j++)
	  {
	    DataType temp=0;
	    for(int i=0;i<eigVecs.rows();i++)
	      {
		int index=rank*partdim;
		if( configs.configs[index+i].test(j) )
		  temp+=pow(eigVecs(i,orb),2);
	      }
	    DataType sum;
	    reduce(diagcom,temp,sum,std::plus<DataType>(),0);
	    if(rank==0)
	      OccuNums[j]=sum;
	  }
      }
  }
  
  void SMSolver::printOccuNums()
  {
    if(world.rank()==0)
      {
	using std::cout;
	using std::setw;
	using std::endl;
	cout<<setw(4)<<"l"<<"\t"<<setw(4)<<"jj"<<"\t"<<setw(4)<<"mm"<<"\t"<<setw(4)<<"tz"<<"\t"<<setw(10)<<"OccuNums"<<endl;

	for(int i=0;i<OccuNums.size();i++)
	  {
	    cout<<setw(4)<<configs.Orbitals[i].l<<"\t"<<setw(4)<<configs.Orbitals[i].jj<<"\t"<<setw(4)<<configs.Orbitals[i].mm<<"\t"<<setw(4)<<configs.Orbitals[i].tz<<"\t"<<setw(10)<<OccuNums[i]<<endl;
	  }
      }
  }



  void SMSolver::printStates(int num)
  {
    if(world.rank()==0)
      {
	using std::cout;
	using std::setw;
	using std::endl;
	if(num==-1||num > numStates) num=numStates;//default num=-1, print all
      
	cout<<"Num"<<"\t"<<setw(15)<<"Energy"<<"\t"<<setw(15)<<"J(J+1)"<<endl;          
	for(int i=0;i<num;i++)
	  {
	    //		cout<<i<<"\t"<<setw(10)<<eigVals(i)<<"\t"<<setw(8)<<int(sqrt(4*J2s(i)+1)-1 + 0.5)/2.<<endl;
	    cout<<i<<"\t"<<setw(15)<<eigVals(i)<<"\t"<<setw(15)<<J2s(i)<<endl;
	    //	cout<<i<<"\t"<<setw(10)<<eigVals(i)<<"\t"<<setw(8)<<int(J2s(i).real() + 0.5)<<endl;
	  }
      }
  }




}
